import java.util.*;

class Solution {
    // 大顶堆 O(nlogk)
    public int[] getLeastNumbers(int[] arr, int k) {
        int[] res = new int[k];
        if (arr == null || arr.length == 0) {
            throw new RuntimeException("invalid parameters");
        }
        if (k == 0) {
            return res;
        }
        PriorityQueue<Integer> queue = new PriorityQueue<>((o1, o2) -> o2 - o1);
        for (int i = 0; i < k; i++) {
            queue.add(arr[i]);
        }
        for (int i = k; i < arr.length; i++) {
            if (arr[i] < queue.peek()) {
                queue.poll();
                queue.add(arr[i]);
            }
        }
        for (int i = 0; i < k; i++) {
            res[i] = queue.poll();
        }
        return res;
    }

    // 小顶堆 O(nlogn)
    public int[] getLeastNumbers1(int[] arr, int k) {
        if (arr == null || arr.length == 0) {
            throw new RuntimeException("invalid parameters");
        }
        PriorityQueue<Integer> queue = new PriorityQueue<>();
        for (int i : arr) {
            queue.add(i);
        }
        int[] res = new int[k];
        for (int i = 0; i < k; i++) {
            res[i] = queue.poll();
        }
        return res;
    }

    // 排序
    public int[] getLeastNumbers3(int[] arr, int k) {
        Arrays.sort(arr);
        return Arrays.copyOf(arr, k);
    }


    int num;

    // 快排思想 O(n)
    public int[] getLastNumber4(int[] arr, int k) {
        if (arr == null || arr.length == 0 || k == 0) return new int[0];
        num = k - 1;
        quickSort(arr, 0, arr.length - 1);
        return Arrays.copyOf(arr, k);
    }

    private void quickSort(int[] arr, int left, int right) {
        int pivot = partition(arr, left, right);
        if (pivot > num) {
            quickSort(arr, left, pivot - 1);
        }
        if (pivot < num) {
            quickSort(arr, pivot + 1, right);
        }
    }

    private int partition(int[] arr, int left, int right) {
        int mid = (left + right) >> 1;
        if (arr[mid] > arr[right]) {
            swap(arr, mid, right);
        }
        if (arr[left] > arr[right]) {
            swap(arr, left, right);
        }
        if (arr[left] < arr[mid]) {
            swap(arr, left, mid);
        }
        int pivot = arr[left];
        while (left < right) {
            while (left < right && arr[right] >= pivot) right--;
            arr[left] = arr[right];
            while (left < right && arr[left] <= pivot) left++;
            arr[right] = arr[left];
        }
        arr[left] = pivot;
        return left;
    }

    private void swap(int[] arr, int i, int j) {
        arr[i] = arr[i] ^ arr[j];
        arr[j] = arr[i] ^ arr[j];
        arr[i] = arr[i] ^ arr[j];
    }

    public static void main(String[] args) {
        Solution s = new Solution();
        System.out.println(Arrays.toString(s.getLastNumber4(new int[]{0, 1, 2, 1}, 2)));
    }
}